import abc
from typing import Literal

from zkl_aiutils_training import SteppingTaskPlugin

from .progress_collector import ProgressCollector


class ActionScheduling(SteppingTaskPlugin, ProgressCollector, abc.ABC):
    def __init__(self, *,
        progress_metric_name: str = 'step',
        action_interval: int | float | None,
        action_mode: Literal['pre', 'post'] = 'post',
    ):
        super().__init__(progress_metric_name=progress_metric_name)
        self.action_mode = action_mode
        self.action_interval = action_interval
        self.action_threshold: int | float | None = None

    def on_before_run(self):
        progress = self._get_progress()
        if self.action_mode == 'pre' and progress == 0:
            self.action_threshold = 0
        else:
            self._update_threshold(progress)
        self._prepare(progress)

    def on_before_step(self):
        if self.action_mode == 'pre':
            self._ensure_after_providers()
            progress = self._get_progress()
            if self._check_threshold(progress):
                self._action(progress)
                self._update_threshold(progress)

    def on_after_step(self):
        if self.action_mode == 'post':
            self._ensure_after_providers()
            progress = self._get_progress()
            if self._check_threshold(progress):
                self._action(progress)
                self._update_threshold(progress)

    def on_after_run(self):
        self._release()
        super().on_after_run()

    # action

    def _update_threshold(self, progress: int | float):
        self.action_threshold = (
            + progress + self.action_interval
            - progress % self.action_interval
        ) if self.action_interval is not None else None

    def _check_threshold(self, progress: int | float) -> bool:
        if self.action_threshold is None:
            return False
        return progress >= self.action_threshold

    def _prepare(self, progress: int | float):
        pass

    @abc.abstractmethod
    def _action(self, progress: int | float):
        pass

    def _release(self):
        pass
